from langchain.chains import LLMChain
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.output_parsers import StructuredOutputParser
from langchain.output_parsers import ResponseSchema
from langchain.tools import tool
from langchain_core.messages import HumanMessage

import os

@tool
def add(a:int, b:int) -> int:
    """Add two numbers together"""
    return a + b

@tool
def multiply(a:int, b:int) -> int:
    """Multiply two numbers together"""
    return a * b

def use_tools() -> dict:
    tools = [add, multiply]
    try:
        api_key = os.environ['API_KEY']
        api_base = os.environ['API_BASE']
        model_name = os.environ['MODEL_NAME']
    except KeyError:
        raise ValueError('DEEPSEEK_API_KEY未在环境变量中配置')
    llm = ChatOpenAI(
        model=model_name,
        openai_api_key=api_key,
        openai_api_base=api_base,
        timeout=1800,
    )
    
    llm_with_tools = llm.bind_tools(tools)
    query = "What is 3 * 12? Also, what is 11 + 49?"
    messages = [HumanMessage(query)]
    ai_msg = llm_with_tools.invoke(messages)
    print(ai_msg.tool_calls)
    messages.append(ai_msg)
    for tool_call in ai_msg.tool_calls:
        selected_tool = {"add": add, "multiply": multiply}[tool_call["name"].lower()]
        tool_msg = selected_tool.invoke(tool_call)
        messages.append(tool_msg)

    print("="*100)
    print(messages)
    print("="*100)
    
    msg = llm_with_tools.invoke(messages) 
    print(msg)
    
    return {
            "aimsg": msg,
        }


